"""
内存优化和垃圾回收服务
"""

import asyncio
import gc
import psutil
import sys
import time
import weakref
import threading
from typing import Dict, List, Any, Optional, Callable, Set
from dataclasses import dataclass, field
from enum import Enum
from collections import defaultdict
from datetime import datetime, timedelta

from ...core.exceptions import BusinessLogicError, ValidationError
from ...core.logging import get_logger
from ...core.config import settings

logger = get_logger(__name__)


class MemoryStatus(Enum):
    """内存状态"""
    HEALTHY = "healthy"         # 健康
    WARNING = "warning"         # 警告
    CRITICAL = "critical"       # 严重
    OUT_OF_MEMORY = "out_of_memory"  # 内存不足


class GCStrategy(Enum):
    """垃圾回收策略"""
    CONSERVATIVE = "conservative"  # 保守策略
    BALANCED = "balanced"          # 平衡策略
    AGGRESSIVE = "aggressive"      # 激进策略
    AUTO = "auto"                  # 自动策略


@dataclass
class MemorySnapshot:
    """内存快照"""
    timestamp: datetime = field(default_factory=datetime.utcnow)
    rss_mb: float = 0.0              # 物理内存 (MB)
    vms_mb: float = 0.0              # 虚拟内存 (MB)
    percent: float = 0.0             # 内存使用百分比
    available_mb: float = 0.0        # 可用内存 (MB)
    gc_counts: Dict[int, int] = field(default_factory=dict)  # GC统计
    
    # 对象统计
    object_count: int = 0
    type_counts: Dict[str, int] = field(default_factory=dict)
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            "timestamp": self.timestamp.isoformat(),
            "rss_mb": self.rss_mb,
            "vms_mb": self.vms_mb,
            "percent": self.percent,
            "available_mb": self.available_mb,
            "gc_counts": self.gc_counts,
            "object_count": self.object_count,
            "type_counts": self.type_counts
        }


@dataclass
class GCMetrics:
    """垃圾回收指标"""
    generation: int
    collected: int
    uncollectable: int
    collections: int
    collection_time: float = 0.0


class MemoryTracker:
    """内存跟踪器"""
    
    def __init__(self, max_snapshots: int = 100):
        self.max_snapshots = max_snapshots
        self.snapshots: List[MemorySnapshot] = []
        
        # 弱引用跟踪
        self.weak_refs: Dict[str, weakref.ref] = {}
        self.object_trackers: Dict[str, Set[weakref.ref]] = defaultdict(set)
        
        # 回调函数
        self.memory_callbacks: List[Callable[[MemorySnapshot], None]] = []
        
        # 统计信息
        self.stats = {
            "total_snapshots": 0,
            "peak_memory_mb": 0.0,
            "peak_percent": 0.0,
            "avg_memory_mb": 0.0,
            "gc_collections": 0
        }
    
    def take_snapshot(self) -> MemorySnapshot:
        """拍摄内存快照"""
        try:
            # 获取进程内存信息
            process = psutil.Process()
            memory_info = process.memory_info()
            memory_percent = process.memory_percent()
            
            # 获取系统内存信息
            system_memory = psutil.virtual_memory()
            
            # 获取GC统计
            gc_stats = gc.get_stats()
            gc_counts = {
                stats['generation']: stats['collections'] 
                for stats in gc_stats
            }
            
            # 创建快照
            snapshot = MemorySnapshot(
                rss_mb=memory_info.rss / 1024 / 1024,
                vms_mb=memory_info.vms / 1024 / 1024,
                percent=memory_percent,
                available_mb=system_memory.available / 1024 / 1024,
                gc_counts=gc_counts
            )
            
            # 对象统计
            try:
                snapshot.object_count = len(gc.get_objects())
                
                # 统计对象类型
                type_counts = {}
                for obj in gc.get_objects():
                    obj_type = type(obj).__name__
                    type_counts[obj_type] = type_counts.get(obj_type, 0) + 1
                
                # 只保留前20个类型
                sorted_types = sorted(type_counts.items(), key=lambda x: x[1], reverse=True)
                snapshot.type_counts = dict(sorted_types[:20])
                
            except Exception as e:
                logger.warning(f"Failed to collect object stats: {e}")
            
            # 添加到历史
            self.snapshots.append(snapshot)
            
            # 限制快照数量
            if len(self.snapshots) > self.max_snapshots:
                self.snapshots.pop(0)
            
            # 更新统计
            self._update_stats(snapshot)
            
            # 通知回调
            for callback in self.memory_callbacks:
                try:
                    callback(snapshot)
                except Exception as e:
                    logger.error(f"Memory callback error: {e}")
            
            return snapshot
            
        except Exception as e:
            logger.error(f"Failed to take memory snapshot: {e}")
            return MemorySnapshot()
    
    def get_latest_snapshot(self) -> Optional[MemorySnapshot]:
        """获取最新快照"""
        return self.snapshots[-1] if self.snapshots else None
    
    def get_snapshots(self, limit: int = None) -> List[MemorySnapshot]:
        """获取快照历史"""
        if limit:
            return self.snapshots[-limit:]
        return self.snapshots
    
    def get_memory_trend(self, minutes: int = 30) -> Dict[str, Any]:
        """获取内存趋势"""
        cutoff_time = datetime.utcnow() - timedelta(minutes=minutes)
        recent_snapshots = [
            s for s in self.snapshots 
            if s.timestamp > cutoff_time
        ]
        
        if not recent_snapshots:
            return {}
        
        rss_values = [s.rss_mb for s in recent_snapshots]
        percent_values = [s.percent for s in recent_snapshots]
        
        return {
            "duration_minutes": minutes,
            "snapshot_count": len(recent_snapshots),
            "rss_mb": {
                "current": rss_values[-1],
                "min": min(rss_values),
                "max": max(rss_values),
                "avg": sum(rss_values) / len(rss_values),
                "trend": rss_values[-1] - rss_values[0] if len(rss_values) > 1 else 0
            },
            "percent": {
                "current": percent_values[-1],
                "min": min(percent_values),
                "max": max(percent_values),
                "avg": sum(percent_values) / len(percent_values)
            }
        }
    
    def track_object(self, obj: Any, name: str = None) -> str:
        """跟踪对象"""
        obj_id = name or f"obj_{id(obj)}"
        
        # 创建弱引用
        weak_ref = weakref.ref(obj, lambda ref: self._on_object_deleted(obj_id))
        self.weak_refs[obj_id] = weak_ref
        
        # 添加到类型跟踪器
        obj_type = type(obj).__name__
        self.object_trackers[obj_type].add(weak_ref)
        
        return obj_id
    
    def untrack_object(self, obj_id: str):
        """取消跟踪对象"""
        if obj_id in self.weak_refs:
            weak_ref = self.weak_refs[obj_id]
            
            # 从类型跟踪器中移除
            for obj_type, refs in self.object_trackers.items():
                if weak_ref in refs:
                    refs.remove(weak_ref)
            
            del self.weak_refs[obj_id]
    
    def get_tracked_objects_count(self) -> Dict[str, int]:
        """获取跟踪对象数量"""
        counts = {}
        
        for obj_type, refs in self.object_trackers.items():
            # 清理已释放的弱引用
            alive_refs = [ref for ref in refs if ref() is not None]
            self.object_trackers[obj_type] = set(alive_refs)
            counts[obj_type] = len(alive_refs)
        
        return counts
    
    def add_memory_callback(self, callback: Callable[[MemorySnapshot], None]):
        """添加内存回调"""
        self.memory_callbacks.append(callback)
    
    def _update_stats(self, snapshot: MemorySnapshot):
        """更新统计信息"""
        self.stats["total_snapshots"] = len(self.snapshots)
        self.stats["peak_memory_mb"] = max(self.stats["peak_memory_mb"], snapshot.rss_mb)
        self.stats["peak_percent"] = max(self.stats["peak_percent"], snapshot.percent)
        
        if self.snapshots:
            self.stats["avg_memory_mb"] = sum(s.rss_mb for s in self.snapshots) / len(self.snapshots)
    
    def _on_object_deleted(self, obj_id: str):
        """对象被删除时的回调"""
        logger.debug(f"Tracked object {obj_id} was deleted")


class GCManager:
    """垃圾回收管理器"""
    
    def __init__(self, strategy: GCStrategy = GCStrategy.AUTO):
        self.strategy = strategy
        self.auto_gc = True
        
        # 配置
        self.gc_thresholds = {
            GCStrategy.CONSERVATIVE: (700, 10, 10),      # (threshold0, threshold1, threshold2)
            GCStrategy.BALANCED: (700, 10, 10),           # 默认值
            GCStrategy.AGGRESSIVE: (500, 5, 5),           # 更激进的阈值
            GCStrategy.AUTO: (700, 10, 10)               # 自动调整
        }
        
        # 统计信息
        self.gc_metrics: Dict[int, GCMetrics] = {
            0: GCMetrics(0, 0, 0, 0),
            1: GCMetrics(1, 0, 0, 0),
            2: GCMetrics(2, 0, 0, 0)
        }
        
        # 回调函数
        self.gc_callbacks: List[Callable[[int, GCMetrics], None]] = []
        
        # 自动调整参数
        self.auto_adjust_enabled = True
        self.last_adjustment_time = 0.0
        self.adjustment_interval = 300  # 5分钟
        
        # 设置初始阈值
        self._set_gc_thresholds()
        
        # 注册回调
        gc.callbacks.append(self._gc_callback)
        
        logger.info(f"GCManager initialized with strategy {strategy.value}")
    
    def set_strategy(self, strategy: GCStrategy):
        """设置垃圾回收策略"""
        self.strategy = strategy
        self._set_gc_thresholds()
        
        logger.info(f"GC strategy changed to {strategy.value}")
    
    def force_gc(self, generation: int = None) -> GCMetrics:
        """强制垃圾回收"""
        start_time = time.time()
        
        try:
            if generation is None:
                # 回收所有代
                collected, uncollectable = gc.collect()
                generation = -1
            else:
                # 回收指定代
                collected, uncollectable = gc.collect(generation)
            
            collection_time = time.time() - start_time
            
            # 更新指标
            if generation in self.gc_metrics:
                metrics = self.gc_metrics[generation]
                metrics.collected += collected
                metrics.uncollectable += uncollectable
                metrics.collections += 1
                metrics.collection_time += collection_time
            
            # 获取GC统计
            stats = gc.get_stats()
            if generation >= 0 and generation < len(stats):
                stats[generation]['collections'] += 1
            
            # 通知回调
            for callback in self.gc_callbacks:
                try:
                    callback(generation, self.gc_metrics[generation])
                except Exception as e:
                    logger.error(f"GC callback error: {e}")
            
            logger.info(f"GC completed for generation {generation}: "
                       f"collected={collected}, uncollectable={uncollectable}, "
                       f"time={collection_time:.3f}s")
            
            return self.gc_metrics.get(generation, GCMetrics(generation, collected, uncollectable, 1, collection_time))
            
        except Exception as e:
            logger.error(f"GC failed: {e}")
            return GCMetrics(generation, 0, 0, 0, 0)
    
    def get_gc_metrics(self) -> Dict[int, GCMetrics]:
        """获取垃圾回收指标"""
        return dict(self.gc_metrics)
    
    def get_gc_stats(self) -> List[Dict[str, Any]]:
        """获取GC统计信息"""
        return gc.get_stats()
    
    def set_gc_thresholds(self, threshold0: int, threshold1: int, threshold2: int):
        """设置GC阈值"""
        gc.set_threshold(threshold0, threshold1, threshold2)
        logger.info(f"GC thresholds set to ({threshold0}, {threshold1}, {threshold2})")
    
    def enable_auto_gc(self):
        """启用自动垃圾回收"""
        self.auto_gc = True
        gc.enable()
        logger.info("Auto GC enabled")
    
    def disable_auto_gc(self):
        """禁用自动垃圾回收"""
        self.auto_gc = False
        gc.disable()
        logger.info("Auto GC disabled")
    
    def optimize_gc(self, memory_pressure: float) -> GCMetrics:
        """根据内存压力优化GC"""
        if memory_pressure > 0.8:
            # 高内存压力，激进回收
            return self.force_gc(0)  # 回收第0代
        elif memory_pressure > 0.6:
            # 中等内存压力
            return self.force_gc(1)  # 回收第1代
        else:
            # 低内存压力，正常回收
            return self.force_gc(2)  # 回收第2代
    
    def adjust_gc_strategy(self, memory_trend: Dict[str, Any]):
        """自动调整GC策略"""
        if not self.auto_adjust_enabled:
            return
        
        current_time = time.time()
        if current_time - self.last_adjustment_time < self.adjustment_interval:
            return
        
        try:
            rss_trend = memory_trend.get("rss_mb", {}).get("trend", 0)
            current_percent = memory_trend.get("percent", {}).get("current", 0)
            
            new_strategy = self.strategy
            
            if self.strategy == GCStrategy.AUTO:
                # 自动策略逻辑
                if rss_trend > 100 or current_percent > 80:  # 内存增长快或使用率高
                    new_strategy = GCStrategy.AGGRESSIVE
                elif current_percent > 60:
                    new_strategy = GCStrategy.BALANCED
                else:
                    new_strategy = GCStrategy.CONSERVATIVE
            
            # 应用新策略
            if new_strategy != self.strategy:
                self.set_strategy(new_strategy)
                self.last_adjustment_time = current_time
                
                logger.info(f"GC strategy auto-adjusted to {new_strategy.value}")
        
        except Exception as e:
            logger.error(f"Failed to adjust GC strategy: {e}")
    
    def add_gc_callback(self, callback: Callable[[int, GCMetrics], None]):
        """添加GC回调"""
        self.gc_callbacks.append(callback)
    
    def _set_gc_thresholds(self):
        """设置GC阈值"""
        thresholds = self.gc_thresholds.get(self.strategy, (700, 10, 10))
        gc.set_threshold(*thresholds)
    
    def _gc_callback(self, phase: str, info: Dict[str, Any]):
        """GC回调函数"""
        try:
            generation = info.get('generation', 0)
            
            # 更新指标
            if generation in self.gc_metrics:
                metrics = self.gc_metrics[generation]
                metrics.collections += 1
                metrics.collected += info.get('collected', 0)
                metrics.uncollectable += info.get('uncollectable', 0)
        
        except Exception as e:
            logger.error(f"GC callback error: {e}")


class MemoryOptimizer:
    """内存优化器"""
    
    def __init__(self):
        self.memory_tracker = MemoryTracker()
        self.gc_manager = GCManager()
        
        # 配置
        self.monitoring_interval = 30  # 30秒
        self.optimization_enabled = True
        self.memory_thresholds = {
            "warning": 70,    # 70%
            "critical": 85,   # 85%
            "emergency": 95   # 95%
        }
        
        # 控制状态
        self.running = False
        self.monitor_task = None
        
        # 统计信息
        self.stats = {
            "optimizations_performed": 0,
            "gc_forced_count": 0,
            "memory_freed_mb": 0.0,
            "peak_memory_saved_mb": 0.0
        }
        
        # 回调函数
        self.optimization_callbacks: List[Callable[[str, Dict[str, Any]], None]] = []
        
        logger.info("MemoryOptimizer initialized")
    
    async def start(self):
        """启动内存优化器"""
        if self.running:
            return
        
        self.running = True
        self.optimization_enabled = True
        
        # 启动监控任务
        self.monitor_task = asyncio.create_task(self._monitoring_loop())
        
        logger.info("MemoryOptimizer started")
    
    async def stop(self):
        """停止内存优化器"""
        if not self.running:
            return
        
        self.running = False
        
        if self.monitor_task:
            self.monitor_task.cancel()
            try:
                await self.monitor_task
            except asyncio.CancelledError:
                pass
        
        logger.info("MemoryOptimizer stopped")
    
    def optimize_now(self, aggressive: bool = False) -> Dict[str, Any]:
        """立即执行内存优化"""
        if not self.optimization_enabled:
            return {"status": "disabled"}
        
        start_time = time.time()
        results = {
            "status": "completed",
            "actions_performed": [],
            "memory_freed_mb": 0.0,
            "time_seconds": 0.0
        }
        
        try:
            # 拍摄优化前快照
            before_snapshot = self.memory_tracker.take_snapshot()
            
            # 执行垃圾回收
            if aggressive:
                # 激进优化：强制所有代的GC
                for gen in range(3):
                    metrics = self.gc_manager.force_gc(gen)
                    results["actions_performed"].append(f"forced_gc_generation_{gen}")
                    self.stats["gc_forced_count"] += 1
            else:
                # 普通优化：根据内存压力选择GC代
                memory_pressure = before_snapshot.percent / 100.0
                metrics = self.gc_manager.optimize_gc(memory_pressure)
                results["actions_performed"].append(f"optimized_gc_pressure_{memory_pressure:.2f}")
                self.stats["gc_forced_count"] += 1
            
            # 清理未使用的缓存
            cache_freed = self._clear_caches()
            if cache_freed > 0:
                results["actions_performed"].append(f"cleared_caches_{cache_freed}mb")
                results["memory_freed_mb"] += cache_freed
            
            # 清理弱引用
            weak_ref_freed = self._cleanup_weak_refs()
            if weak_ref_freed > 0:
                results["actions_performed"].append(f"cleaned_weak_refs_{weak_ref_freed}")
            
            # 压缩数据结构
            self._compact_data_structures()
            results["actions_performed"].append("compacted_structures")
            
            # 拍摄优化后快照
            after_snapshot = self.memory_tracker.take_snapshot()
            
            # 计算释放的内存
            memory_freed = before_snapshot.rss_mb - after_snapshot.rss_mb
            results["memory_freed_mb"] += memory_freed
            
            # 更新统计
            self.stats["optimizations_performed"] += 1
            self.stats["memory_freed_mb"] += results["memory_freed_mb"]
            
            if memory_freed > 0:
                logger.info(f"Memory optimization freed {memory_freed:.2f} MB")
            else:
                logger.info("Memory optimization completed (no significant memory freed)")
            
            # 通知回调
            await self._notify_optimization("manual_optimization", results)
            
        except Exception as e:
            results["status"] = "error"
            results["error"] = str(e)
            logger.error(f"Memory optimization failed: {e}")
        
        finally:
            results["time_seconds"] = time.time() - start_time
        
        return results
    
    def get_memory_status(self) -> MemoryStatus:
        """获取内存状态"""
        latest = self.memory_tracker.get_latest_snapshot()
        if not latest:
            return MemoryStatus.HEALTHY
        
        if latest.percent >= self.memory_thresholds["emergency"]:
            return MemoryStatus.OUT_OF_MEMORY
        elif latest.percent >= self.memory_thresholds["critical"]:
            return MemoryStatus.CRITICAL
        elif latest.percent >= self.memory_thresholds["warning"]:
            return MemoryStatus.WARNING
        else:
            return MemoryStatus.HEALTHY
    
    def get_optimization_report(self) -> Dict[str, Any]:
        """获取优化报告"""
        latest = self.memory_tracker.get_latest_snapshot()
        memory_trend = self.memory_tracker.get_memory_trend(minutes=60)
        gc_metrics = self.gc_manager.get_gc_metrics()
        tracked_objects = self.memory_tracker.get_tracked_objects_count()
        
        return {
            "memory_status": self.get_memory_status().value,
            "current_memory": latest.to_dict() if latest else None,
            "memory_trend": memory_trend,
            "gc_metrics": {
                gen: {
                    "collected": metrics.collected,
                    "collections": metrics.collections,
                    "collection_time": metrics.collection_time,
                    "avg_collection_time": metrics.collection_time / max(1, metrics.collections)
                }
                for gen, metrics in gc_metrics.items()
            },
            "tracked_objects": tracked_objects,
            "optimization_stats": dict(self.stats),
            "gc_strategy": self.gc_manager.strategy.value,
            "optimization_enabled": self.optimization_enabled
        }
    
    def add_optimization_callback(self, callback: Callable[[str, Dict[str, Any]], None]):
        """添加优化回调"""
        self.optimization_callbacks.append(callback)
    
    async def _monitoring_loop(self):
        """监控循环"""
        while self.running:
            try:
                # 拍摄内存快照
                snapshot = self.memory_tracker.take_snapshot()
                
                # 检查是否需要优化
                if self.optimization_enabled:
                    await self._check_and_optimize(snapshot)
                
                # 等待下一次检查
                await asyncio.sleep(self.monitoring_interval)
                
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Memory monitoring error: {e}")
                await asyncio.sleep(10)
    
    async def _check_and_optimize(self, snapshot: MemorySnapshot):
        """检查并执行优化"""
        memory_status = self.get_memory_status()
        
        if memory_status == MemoryStatus.CRITICAL:
            # 严重状态：立即执行优化
            logger.warning("Critical memory usage, performing emergency optimization")
            await self._emergency_optimization()
        
        elif memory_status == MemoryStatus.OUT_OF_MEMORY:
            # 内存不足：激进优化
            logger.error("Out of memory condition, performing aggressive optimization")
            await self._aggressive_optimization()
    
    async def _emergency_optimization(self):
        """紧急优化"""
        results = self.optimize_now(aggressive=False)
        await self._notify_optimization("emergency_optimization", results)
    
    async def _aggressive_optimization(self):
        """激进优化"""
        results = self.optimize_now(aggressive=True)
        await self._notify_optimization("aggressive_optimization", results)
    
    def _clear_caches(self) -> float:
        """清理缓存"""
        freed_mb = 0.0
        
        try:
            # 清理模块缓存
            import sys
            before_modules = len(sys.modules)
            
            # 清理未使用的模块
            modules_to_remove = []
            for name, module in sys.modules.items():
                if not name.startswith('_') and not name.startswith('sys'):
                    modules_to_remove.append(name)
            
            # 保守地只清理明显未使用的模块
            # 这里简化处理，实际中需要更复杂的逻辑
            
            # 计算释放的内存（估算）
            freed_mb = len(modules_to_remove) * 0.5  # 估算每个模块0.5MB
            
            logger.debug(f"Cache cleanup freed ~{freed_mb:.2f} MB")
            
        except Exception as e:
            logger.error(f"Cache cleanup failed: {e}")
        
        return freed_mb
    
    def _cleanup_weak_refs(self) -> int:
        """清理弱引用"""
        cleaned = 0
        
        try:
            for obj_type, refs in self.memory_tracker.object_trackers.items():
                # 清理已释放的弱引用
                alive_refs = [ref for ref in refs if ref() is not None]
                cleaned += len(refs) - len(alive_refs)
                self.memory_tracker.object_trackers[obj_type] = set(alive_refs)
            
            if cleaned > 0:
                logger.debug(f"Cleaned up {cleaned} dead weak references")
                
        except Exception as e:
            logger.error(f"Weak reference cleanup failed: {e}")
        
        return cleaned
    
    def _compact_data_structures(self):
        """压缩数据结构"""
        try:
            # 强制Python内部数据结构压缩
            # 这是一个简化的实现，实际中可能需要更复杂的逻辑
            
            # 调用GC来压缩对象
            gc.collect()
            
            logger.debug("Data structure compaction completed")
            
        except Exception as e:
            logger.error(f"Data structure compaction failed: {e}")
    
    async def _notify_optimization(self, optimization_type: str, results: Dict[str, Any]):
        """通知优化事件"""
        for callback in self.optimization_callbacks:
            try:
                if asyncio.iscoroutinefunction(callback):
                    await callback(optimization_type, results)
                else:
                    callback(optimization_type, results)
            except Exception as e:
                logger.error(f"Optimization callback error: {e}")


# 全局内存优化器实例
memory_optimizer = MemoryOptimizer()
